import pandas as pd
from datetime import datetime
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import matplotlib.lines as mlines
import seaborn as sns
import numpy as np
from string import punctuation

import warnings
warnings.filterwarnings('ignore')

c = 1000

label_fontsize = 30
tick_fontsize = 25
markersize = 18

# Import date field as date type
dateparse = lambda x: datetime.strptime(x, '%Y-%m-%d')

df = pd.read_csv('../data/gpo_final_data/narratives_complete_with_metadata_manual_labels_rich_{0}.csv'.format(c), parse_dates=['date'], date_parser=dateparse)

required_fields = ['sentence_raw', 'gpo_id',
                   'narrative', 'date', 'congress', 'frequency', 'sentiment_neg',
                   'sentiment_pos', 'sentiment_neu', 'sentiment_compound', 'or', 'log_or']

main_data = df[required_fields]

del df, required_fields

main_data['year'] = main_data['date'].dt.year

_top_narratives_df = main_data[main_data['narrative'].str.contains('god bless america|god bless troop|saddam hussein have weapon mass destruction|saddam hussein pose threat')]
_top_narratives_df['n'] = 1

_top_narratives_grouped = _top_narratives_df.groupby(['year', 'narrative']).count()['n'].unstack()

# Fill NaN values with zero
_top_narratives_grouped = _top_narratives_grouped.fillna(0)

# Setting labels for plot legend
_legends = list(_top_narratives_grouped)

# Plotting
_fig, _ax = plt.subplots(figsize=(35,10))
_ax.xaxis.set_major_locator(ticker.MultipleLocator(1))
_ax.xaxis.label.set_fontsize(label_fontsize)
_ax.set_xlabel('\n Year')
_ax.yaxis.label.set_fontsize(label_fontsize)
_ax.set_ylabel('Count \n')
_ax.set_facecolor((.18, .31, .31, 0.0))
_ax.grid(color='#dfede2')


_g = sns.lineplot(data=_top_narratives_grouped,
                 markers = True,
                 dashes = False,
                 ax=_ax,
                 markersize=markersize,
                 linewidth=2.5)
_ax.set_prop_cycle(None)
_top_narratives_grouped.plot.area(stacked=False, alpha=0.5, ax=_ax)
plt.legend(title='Narratives', loc='upper right', labels=_legends, title_fontsize=label_fontsize, fontsize=label_fontsize)
plt.xticks([1995, 2000, 2005, 2010, 2015], fontsize=tick_fontsize)
plt.yticks(fontsize=tick_fontsize)
_ax.set_xlabel('\n Year')
#plt.show(_g)
plt.savefig('../figures/Figure_3_a.pdf')
